{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Train Set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating Atom Models with Different Crystal Structures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.atoms import SymbolInfo, Morph\n",
    "import numpy as np\n",
    "\n",
    "train_set = []\n",
    "SymbolInfos = [\n",
    "   SymbolInfo('MoTaVW', 'bcc', 3.32),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generating and Saving Deformed Atom Models with Random Strain and Displacement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.frames import MultiMol\n",
    "\n",
    "frames = []\n",
    "for SymbolInfo in SymbolInfos:\n",
    "    atoms = SymbolInfo.create_bulk_atoms((2, 2, 2))\n",
    "    frames.append(atoms)\n",
    "\n",
    "deform_scale = np.arange(0.95, 1.06, 0.05)\n",
    "strain_ratio = 0.04\n",
    "max_displacement = 0.4\n",
    "\n",
    "init_1 = MultiMol(frames).deform(deform_scale)\n",
    "init_2 = MultiMol(init_1).random_strain(strain_ratio)\n",
    "init_3 = MultiMol(init_1).random_displacement(max_displacement)\n",
    "\n",
    "init = init_1 + init_2 + init_3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generating Atom Models with Various Types of Defects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase import Atoms\n",
    "\n",
    "dimers = []\n",
    "distances = [1.1 + i * 0.1 for i in range(10)]\n",
    "for symbol_info in SymbolInfos:\n",
    "    symbol1 = symbol_info.symbols[0]\n",
    "    symbol2 = symbol_info.symbols[-1]\n",
    "    for distance in distances:\n",
    "        dimer = Atoms(symbols= [symbol1, symbol2], positions=[(0, 0, 0), (0, 0, distance)], pbc = [True, True, True] ,cell=(20, 30, 40))\n",
    "        dimers.append(dimer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mono_vacancy = []\n",
    "for symbol_info in SymbolInfos:\n",
    "    atoms = symbol_info.create_bulk_atoms((3, 4, 5))\n",
    "    Morph(atoms).create_vacancy()\n",
    "    mono_vacancy.append(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nths = [1,2,3,4,5]\n",
    "di_vacancies = []\n",
    "for symbol_info in SymbolInfos:\n",
    "    for nth in nths:\n",
    "        atoms = symbol_info.create_bulk_atoms((3, 4, 5)) \n",
    "        Morph(atoms).create_divacancies(nth)\n",
    "        di_vacancies.append(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vacancies = []\n",
    "for symbol_info in SymbolInfos:\n",
    "    for n in range(3, 10, 3):\n",
    "        atoms = symbol_info.create_bulk_atoms((3, 4, 5)) \n",
    "        Morph(atoms).create_vacancies(n)\n",
    "        vacancies.append(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sia = []\n",
    "vectors = [(1,1,1),(1,0,0),(1,1,0)]\n",
    "for symbol_info in SymbolInfos:\n",
    "    for vector in vectors:\n",
    "        atoms = symbol_info.create_bulk_atoms((3,4,5)) \n",
    "        Morph(atoms).create_self_interstitial_atom(vector)\n",
    "        sia.append(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectors = [(1,1,1),(1,0,0),(1,1,0)]\n",
    "nths = [1,2,3,4,5]\n",
    "di_sias = []\n",
    "for symbol_info in SymbolInfos:\n",
    "    for vector1 in vectors:\n",
    "        for vector2 in vectors:\n",
    "            for nth in nths:\n",
    "                atoms = symbol_info.create_bulk_atoms((3,4,5))\n",
    "                Morph(atoms).create_di_self_interstitial_atoms(vector1=vector1, vector2=vector2, nth=nth)\n",
    "                di_sias.append(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fks = []\n",
    "for symbol_info in SymbolInfos:\n",
    "    atoms = symbol_info.create_bulk_atoms((3,4,5))\n",
    "    Morph(atoms).create_fks(10)\n",
    "    fks.append(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.build import surface\n",
    "\n",
    "millers = [(1,1,0),(0,0,1),(1,1,1),(1,1,2)]\n",
    "surf = []\n",
    "for symbol_info in SymbolInfos:\n",
    "    for miller in millers:\n",
    "        atoms = symbol_info.create_bulk_atoms()\n",
    "        slab = surface(atoms, miller, layers = 10, vacuum=10) *(2, 2, 1)\n",
    "        surf.append(slab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exploring Structures Using GPUMD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.io import read_xyz\n",
    "from wizard.frames import MultiMol\n",
    "\n",
    "Temp_Struts = []\n",
    "temperatures = [50, 300, 800, 1300, 1700, 2300, 3000, 4000, 5000]\n",
    "for symbol_info in SymbolInfos:\n",
    "    for temperature in temperatures:\n",
    "        atoms = symbol_info.create_bulk_atoms((3,3,3))\n",
    "        dirname = f'{symbol_info.formula}/{symbol_info.structure}/{temperature}K/relax'\n",
    "        run_in=['potential nep.txt',\n",
    "                f'velocity {temperature}',   \n",
    "                'time_step 1',\n",
    "                f'ensemble npt_mttk temp {temperature} {temperature} iso 0 0',\n",
    "                'dump_thermo 10000',\n",
    "                'dump_exyz 200000',\n",
    "                'dump_restart 1000000',\n",
    "                'run 1000000']\n",
    "        Morph(atoms).gpumd(dirname=dirname, run_in=run_in)\n",
    "        Temp_Struts += read_xyz(dirname + '/dump.xyz')\n",
    "\n",
    "deform_scale = np.arange(0.95, 1.06, 0.05)\n",
    "strain_ratio = 0.05\n",
    "Temp_Struts_Deform = MultiMol(Temp_Struts).deform(deform_scale)\n",
    "Temp_Struts_Strain = MultiMol(Temp_Struts).random_strain(strain_ratio)\n",
    "\n",
    "Temp_Struts = Temp_Struts_Deform + Temp_Struts_Strain\n",
    "MultiMol(Temp_Struts).dump('Temp_Struts.xyz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.io import read_xyz, read_restart\n",
    "from wizard.frames import MultiMol\n",
    "\n",
    "temperature = 300 #K\n",
    "strain_rate = 2e8\n",
    "dt = 1e-15 #fs\n",
    "utc_tensile = []\n",
    "for symbol_info in SymbolInfos:\n",
    "    atoms = read_restart(f'{symbol_info.symbol}/{symbol_info.structure}/{temperature}K/relax/restart.xyz')\n",
    "    dirname = f'{symbol_info.symbol}/{symbol_info.structure}/{temperature}K/utc_tensile'\n",
    "    length = atoms.cell[2, 2]\n",
    "    strain = strain_rate * dt * length\n",
    "    run_in = ['potential nep.txt', \n",
    "              'velocity 300', \n",
    "              'time_step 1',\n",
    "              f'ensemble npt_scr {temperature} {temperature} 100 0 0 0 100 100 100 1000',\n",
    "              f'deform {strain} 0 0 1', \n",
    "              'dump_thermo 1000', \n",
    "              'dump_exyz 200000', \n",
    "              'dump_restart 10000',\n",
    "              'run 2000000']\n",
    "    Morph(atoms).gpumd(dirname=dirname, run_in=run_in)\n",
    "    utc_tensile += read_xyz(dirname + '/dump.xyz') \n",
    "\n",
    "MultiMol(utc_tensile).dump('utc_tensile.xyz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.io import read_xyz\n",
    "from wizard.frames import MultiMol\n",
    "\n",
    "frames = read_xyz('defect.xyz')\n",
    "Temp_Struts = []\n",
    "temperatures = [50, 300, 800, 1300, 1700, 2300, 3000, 4000, 5000]\n",
    "for i, atoms in enumerate(frames):\n",
    "    for temperature in temperatures:\n",
    "        dirname = f'defects/{i}/{temperature}K/relax'\n",
    "        run_in=['potential nep.txt',\n",
    "                f'velocity {temperature}',   \n",
    "                'time_step 1',\n",
    "                f'ensemble npt_mttk temp {temperature} {temperature} iso 0 0',\n",
    "                'dump_thermo 10000',\n",
    "                'dump_exyz 200000',\n",
    "                'dump_restart 1000000',\n",
    "                'run 1000000']\n",
    "        Morph(atoms).gpumd(dirname=dirname, run_in=run_in)\n",
    "        Temp_Struts += read_xyz(dirname + '/dump.xyz')\n",
    "\n",
    "MultiMol(Temp_Struts).dump('Defect_Temp_Struts.xyz')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
